package org.taj.adapter.llama;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.annotation.JSONField;
import okhttp3.*;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import okhttp3.sse.EventSources;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

public class LLaMARequest {
    public static class Message {
        @JSONField(name = "id")
        private long id;
        @JSONField(name = "role")
        private String role;
        @JSONField(name = "content")
        private String content;

        public Message(long id, String role, String content) {
            this.id = id;
            this.role = role;
            this.content = content;
        }

        public long getId() {
            return id;
        }

        public String getRole() {
            return role;
        }

        public String getContent() {
            return content;
        }

        public void setContent(String content) {
            this.content = content;
        }
    }

    public static class Tool {
        @JSONField(name = "type")
        private String type;
        @JSONField(name = "function")
        private Function function;

        public String getType() {
            return type;
        }

        public void setType(String type) {
            this.type = type;
        }

        public Function getFunction() {
            return function;
        }

        public void setFunction(Function function) {
            this.function = function;
        }
    }

    public static class Function {
        @JSONField(name = "name")
        private String name;
        @JSONField(name = "description")
        private String description;
        @JSONField(name = "parameters")
        private Parameter parameters;

        public String getName() {
            return name;
        }

        public void setName(String name) {
            this.name = name;
        }

        public String getDescription() {
            return description;
        }

        public void setDescription(String description) {
            this.description = description;
        }

        public Parameter getParameters() {
            return parameters;
        }

        public void setParameters(Parameter parameters) {
            this.parameters = parameters;
        }
    }

    public static class Parameter {
        @JSONField(name = "type")
        private String type;
        @JSONField(name = "required")
        private List<String> required;
        @JSONField(name = "additionalProperties")
        private boolean additionalProperties;
        @JSONField(name = "properties")
        private Map<String, PropertyItem> properties = new HashMap<>();

        public String getType() {
            return type;
        }

        public void setType(String type) {
            this.type = type;
        }

        public List<String> getRequired() {
            return required;
        }

        public void setRequired(List<String> required) {
            this.required = required;
        }

        public boolean isAdditionalProperties() {
            return additionalProperties;
        }

        public void setAdditionalProperties(boolean additionalProperties) {
            this.additionalProperties = additionalProperties;
        }

        public Map<String, PropertyItem> getProperties() {
            return properties;
        }

        public void setProperties(Map<String, PropertyItem> properties) {
            this.properties = properties;
        }

        public void addProperty(String key, PropertyItem propertyItem) {
            this.properties.put(key, propertyItem);
        }
    }

    public static class PropertyItem {
        @JSONField(name = "type")
        private String type;
        @JSONField(name = "description")
        private String description;

        public String getType() {
            return type;
        }

        public void setType(String type) {
            this.type = type;
        }

        public String getDescription() {
            return description;
        }

        public void setDescription(String description) {
            this.description = description;
        }
    }

    @JSONField(name = "messages")
    private List<Message> messages;
    @JSONField(name = "tools")
    private List<Tool> tools;
    @JSONField(name = "stream")
    private boolean stream;
    @JSONField(name = "cache_prompt")
    private boolean cachePrompt;
    @JSONField(name = "samplers")
    private String samplers;
    @JSONField(name = "temperature")
    private double temperature;
    @JSONField(name = "dynatemp_range")
    private double dynatempRange;
    @JSONField(name = "dynatemp_exponent")
    private double dynatempExponent;
    @JSONField(name = "top_k")
    private int topK;
    @JSONField(name = "top_p")
    private double topP;
    @JSONField(name = "min_p")
    private double minP;
    @JSONField(name = "typical_p")
    private int typicalP;
    @JSONField(name = "xtc_probability")
    private int xtcProbability;
    @JSONField(name = "xtc_threshold")
    private double xtcThreshold;
    @JSONField(name = "repeat_last_n")
    private int repeatLastN;
    @JSONField(name = "repeat_penalty")
    private int repeatPenalty;
    @JSONField(name = "presence_penalty")
    private int presencePenalty;
    @JSONField(name = "frequency_penalty")
    private int frequencyPenalty;
    @JSONField(name = "dry_multiplier")
    private int dryMultiplier;
    @JSONField(name = "dry_base")
    private double dryBase;
    @JSONField(name = "dry_allowed_length")
    private int dryAllowedLength;
    @JSONField(name = "dry_penalty_last_n")
    private int dryPenaltyLastN;
    @JSONField(name = "max_tokens")
    private int maxTokens;
    @JSONField(name = "timings_per_token")
    private boolean timingsPerToken;

    public LLaMARequest(List<Message> messages) {
        this(messages, null, true, true, "edkypmxt", 0.8, 0.0, 1.0, 40, 0.95,
                0.05, 1, 0, 0.1, 64, 1, 0, 0, 0,
                1.75, 2, -1, -1, false);
    }

    public LLaMARequest(List<Message> messages, List<Tool> tools) {
        this(messages, tools, true, true, "edkypmxt", 0.8, 0.0, 1.0, 40, 0.95,
                0.05, 1, 0, 0.1, 64, 1, 0, 0, 0,
                1.75, 2, -1, -1, false);
        this.stream = false;
    }

    public LLaMARequest(List<Message> messages, List<Tool> tools,
                        boolean stream, boolean cachePrompt, String samplers, double temperature,
                        double dynatempRange, double dynatempExponent,
                        int topK, double topP, double minP, int typicalP, int xtcProbability,
                        double xtcThreshold, int repeatLastN, int repeatPenalty, int presencePenalty,
                        int frequencyPenalty, int dryMultiplier, double dryBase, int dryAllowedLength,
                        int dryPenaltyLastN, int maxTokens, boolean timingsPerToken) {
        this.stream = stream;
        this.cachePrompt = cachePrompt;
        this.samplers = samplers;
        this.temperature = temperature;
        this.dynatempRange = dynatempRange;
        this.dynatempExponent = dynatempExponent;
        this.topK = topK;
        this.topP = topP;
        this.minP = minP;
        this.typicalP = typicalP;
        this.xtcProbability = xtcProbability;
        this.xtcThreshold = xtcThreshold;
        this.repeatLastN = repeatLastN;
        this.repeatPenalty = repeatPenalty;
        this.presencePenalty = presencePenalty;
        this.frequencyPenalty = frequencyPenalty;
        this.dryMultiplier = dryMultiplier;
        this.dryBase = dryBase;
        this.dryAllowedLength = dryAllowedLength;
        this.dryPenaltyLastN = dryPenaltyLastN;
        this.maxTokens = maxTokens;
        this.timingsPerToken = timingsPerToken;
        this.messages = messages;
        this.tools = tools;
    }

    public List<Message> getMessages() {
        return messages;
    }

    public List<Tool> getTools() {
        return tools;
    }

    public boolean isStream() {
        return stream;
    }

    public boolean isCachePrompt() {
        return cachePrompt;
    }

    public String getSamplers() {
        return samplers;
    }

    public double getTemperature() {
        return temperature;
    }

    public double getDynatempRange() {
        return dynatempRange;
    }

    public double getDynatempExponent() {
        return dynatempExponent;
    }

    public int getTopK() {
        return topK;
    }

    public double getTopP() {
        return topP;
    }

    public double getMinP() {
        return minP;
    }

    public int getTypicalP() {
        return typicalP;
    }

    public int getXtcProbability() {
        return xtcProbability;
    }

    public double getXtcThreshold() {
        return xtcThreshold;
    }

    public int getRepeatLastN() {
        return repeatLastN;
    }

    public int getRepeatPenalty() {
        return repeatPenalty;
    }

    public int getPresencePenalty() {
        return presencePenalty;
    }

    public int getFrequencyPenalty() {
        return frequencyPenalty;
    }

    public int getDryMultiplier() {
        return dryMultiplier;
    }

    public double getDryBase() {
        return dryBase;
    }

    public int getDryAllowedLength() {
        return dryAllowedLength;
    }

    public int getDryPenaltyLastN() {
        return dryPenaltyLastN;
    }

    public int getMaxTokens() {
        return maxTokens;
    }

    public boolean isTimingsPerToken() {
        return timingsPerToken;
    }

    public static void main(String[] args) throws InterruptedException {
        List<Tool> tools = new ArrayList<>();
        Tool tool = new Tool();
        tool.setType("function");
        Function function = new Function();
        function.setName("get_current_weather");
        function.setDescription("Get the current weather in a given location");
        Parameter parameter = new Parameter();
        parameter.setType("object");
        parameter.setRequired(List.of("location", "date"));
        PropertyItem location = new PropertyItem();
        location.setType("string");
        location.setDescription("The city and state, e.g. San Francisco, CA");
        PropertyItem date = new PropertyItem();
        date.setType("string");
        date.setDescription("The date, in YYYY-MM-DD format");
        parameter.addProperty("location", location);
        parameter.addProperty("date", date);
        function.setParameters(parameter);
        tool.setFunction(function);
        tools.add(tool);

        LLaMARequest request = new LLaMARequest(
                List.of(new Message(0, "user", "南京2025年3月5日的是否下雨?")), tools
        );

        String json = JSON.toJSONString(request);
        RequestBody body = RequestBody.create(json, MediaType.parse("application/json"));
        Request req = new Request.Builder().url("http://192.168.2.246:8080/v1/chat/completions").post(body).build();
        OkHttpClient client = new OkHttpClient();
        EventSource.Factory factory = EventSources.createFactory(client);
        factory.newEventSource(req, new EventSourceListener() {
            @Override
            public void onOpen(EventSource eventSource, Response response) {
                System.out.println("SSE连接已开启");
            }

            @Override
            public void onEvent(EventSource eventSource, String id, String type, String data) {
                System.out.println("SSE数据: " + data);
            }

            @Override
            public void onClosed(EventSource eventSource) {
                System.out.println("SSE连接已关闭");
            }

            @Override
            public void onFailure(EventSource eventSource, Throwable t, Response response) {
                System.err.println("SSE连接失败: " + t.getMessage());
                if (response != null) {
                    System.err.println("Response Code: " + response.code());
                }
            }
        });

        TimeUnit.SECONDS.sleep(100);
    }
}
